import os
import csv
import numpy as np

import math


def out(name, out_pth, D, movie_id_listall, movie_num, user_id_listall, user_num):
    # 评分矩阵输出
    fp_output = open(
        os.path.join(out_pth, name), mode="w", encoding="utf-8", newline=""
    )
    writer = csv.writer(fp_output)
    writer.writerow([""] + movie_id_listall)
    p = []
    for i in range(user_num):
        p.append(user_id_listall[i])
        for j in range(movie_num):
            p.append(D[i][j])
        writer.writerow(p)
        p = []

    print("输出" + name)


def stage3(data_pth, out_pth):
    fp_score = open(os.path.join(data_pth, "Movie_score.csv"), mode="r", encoding="gbk")

    score_csv = csv.reader(fp_score)
    next(score_csv)
    user_id_listall = []
    movie_id_listall = []
    for row in score_csv:
        if row[0] not in user_id_listall:
            user_id_listall.append(row[0])
        if row[1] not in movie_id_listall:
            movie_id_listall.append(row[1])

    user_num = len(user_id_listall)
    movie_num = len(movie_id_listall)
    user_id_dictall = dict()
    movie_id_dictall = dict()
    for i in range(user_num):
        user_id_dictall.update({user_id_listall[i]: i})
    for i in range(movie_num):
        movie_id_dictall.update({movie_id_listall[i]: i})

    fp_score.close()
    print("D?")
    D = np.zeros((user_num, movie_num), dtype="float")
    # 矩阵维度
    print(D.shape)
    # 修改此项可以修改已知矩阵信息
    fp_train = open(os.path.join(data_pth, "train.csv"), mode="r", encoding="utf-8")
    train_csv = csv.reader(fp_train)
    next(train_csv)
    # 由于给的评分有0分，将所有评分上调一分，矩阵中零为未评分
    for row in train_csv:
        D[user_id_dictall[row[0]]][movie_id_dictall[row[1]]] = int(row[2]) + 1

    out(
        "pingfen.csv",
        out_pth,
        D,
        movie_id_listall,
        movie_num,
        user_id_listall,
        user_num,
    )
    # 统计每个用户平均分
    avg_user = []
    movie_count = []
    sum = 0
    count = 0
    for i in range(user_num):
        for j in range(movie_num):
            if D[i][j] != 0:
                count += 1
                sum += D[i][j]
        movie_count.append(count)
        avg_user.append(sum / count)
        sum = 0
        count = 0
    # print(avg_user[:20])
    # print(movie_count[:5])
    # 矩阵减去平均分
    for i in range(user_num):
        for j in range(movie_num):
            if D[i][j] != 0:
                D[i][j] -= avg_user[i]
    out(
        "quzhongxin.csv",
        out_pth,
        D,
        movie_id_listall,
        movie_num,
        user_id_listall,
        user_num,
    )
    # sim下面根号那坨
    Sqrt = []
    sum = 0
    for i in range(user_num):
        for j in range(movie_num):
            sum += math.pow(D[i][j], 2)
        Sqrt.append(math.sqrt(sum))
        sum = 0
    print(Sqrt[:20])
    # 相似性矩阵，存储任意两个用户的相似性，初始化零
    sum = 0
    SIM = np.zeros((user_num, user_num), dtype="float64")
    for i in range(user_num - 1):
        for j in range(i + 1, user_num):
            for k in range(movie_num):
                if D[i][k] != 0 and D[j][k] != 0:
                    sum += D[i][k] * D[j][k]
            if Sqrt[i] * Sqrt[j] == 0:
                SIM[i][j] = 0
            else:
                SIM[i][j] = sum / (Sqrt[i] * Sqrt[j])
                SIM[j][i] = SIM[i][j]
            sum = 0
    out("sim.csv", out_pth, SIM, user_id_listall, user_num, user_id_listall, user_num)
    k = math.pow(math.log(user_num), 2)
    # 每个user选择k个来预测
    # flag存每个user的k个最相似用户，是记为1
    flag = []
    for i in range(user_num):
        flagflag = [0 for i in range(user_num)]
        max = 0
        maxi = -1
        j = 0
        while j < k:
            for t in range(user_num):
                if flagflag[t] == 1:  # 已经被选出，跳过
                    continue
                else:
                    if SIM[i][t] > max:
                        max = SIM[i][t]
                        maxi = t
            if maxi == -1:
                break
            else:
                flagflag[maxi] = 1
                j = j + 1
            max = 0
            maxi = -1
        flag.append(flagflag)
        if i < 5:
            print(flagflag)
    print(flag[0])

    # 预测
    fp_test = open(os.path.join(data_pth, "test.csv"), mode="r", encoding="utf-8")
    test_csv = csv.reader(fp_test)
    next(test_csv)
    D_true = np.zeros((user_num, movie_num), dtype="float")
    D_predict = np.zeros((user_num, movie_num), dtype="float")
    countrow = 0
    for row in test_csv:
        print(countrow, end=" ")
        countrow += 1
        D_true[user_id_dictall[row[0]]][movie_id_dictall[row[1]]] = int(row[2]) + 1

        simsum = 0
        sum = 0
        for i in range(user_num):
            if (
                flag[user_id_dictall[row[0]]][i] == 1
                and D[i][movie_id_dictall[row[1]]] != 0
            ):
                simsum += SIM[user_id_dictall[row[0]]][i]
                sum += SIM[user_id_dictall[row[0]]][i] * D[i][movie_id_dictall[row[1]]]
        print(sum, end=" ")
        print(simsum)
        if simsum == 0:
            ans = 0
        else:
            ans = sum / simsum

        D_predict[user_id_dictall[row[0]]][movie_id_dictall[row[1]]] = (
            ans + avg_user[user_id_dictall[row[0]]]
        )
    out(
        "true.csv",
        out_pth,
        D_true,
        movie_id_listall,
        movie_num,
        user_id_listall,
        user_num,
    )
    out(
        "predict.csv",
        out_pth,
        D_predict,
        movie_id_listall,
        movie_num,
        user_id_listall,
        user_num,
    )

    # 计算NDCG
    NDCG = []
    DCG = []
    iDCG = []
    # 计算DCG
    for i in range(user_num):
        max = 0
        maxj = -1
        count = 0
        dcg = 0
        flag = [0 for i in range(movie_num)]
        while count < movie_num:
            for j in range(movie_num):
                if flag[j] == 1:
                    continue
                elif D_predict[i][j] > max:
                    max = D_predict[i][j]
                    maxj = j
            if maxj == -1:  # 算完退出
                break
            else:
                dcg += D_true[i][maxj] / math.log2(count + 2)
                flag[j] = 1
                count += 1

            max = 0
            maxj = -1
        DCG.append(dcg)
        print("DCG" + str(i))
    # 计算iDCG
    for i in range(user_num):
        max = 0
        maxj = -1
        count = 0
        idcg = 0
        flag = [0 for i in range(movie_num)]
        while count < movie_num:
            for j in range(movie_num):
                if flag[j] == 1:
                    continue
                elif D_true[i][j] > max:
                    max = D_true[i][j]
                    maxj = j
            if maxj == -1:  # 算完退出
                break
            else:
                idcg += D_true[i][maxj] / math.log2(count + 2)
                flag[j] = 1
                count += 1
            max = 0
            maxj = -1
        iDCG.append(idcg)
        print("iDCG" + str(i))

    print(DCG)
    print(iDCG)
    for i in range(user_num):
        
        NDCG.append(DCG[i] / iDCG[i])

    print(NDCG)

    return


if __name__ == "__main__":
    pth = os.path.split(os.path.realpath(__file__))[0]
    data_pth = os.path.join(pth, "data")
    out_pth = os.path.join(pth, "output")

    stage3(data_pth, out_pth)
